package cn.pomelo.ws.listener;

import cn.pomelo.ws.endpoint.WebSocketServer;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.unfbx.chatgpt.entity.chat.ChatCompletionResponse;
import lombok.NoArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.ResponseBody;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import javax.websocket.Session;

@Slf4j
@NoArgsConstructor
public class OpenAiWebSocketEventSourceListener extends EventSourceListener {
    private Session session;

    public OpenAiWebSocketEventSourceListener(Session session) {
        this.session = session;
    }

    @Override
    public void onOpen(@NotNull EventSource eventSource, @NotNull Response response) {
        log.info("OpenAI建立连接成功...");
    }

    @SneakyThrows
    @Override
    public void onEvent(@NotNull EventSource eventSource, @Nullable String id, @Nullable String type, @NotNull String data) {
        log.info("OpenAI返回数据{}", data);
        if (data.equals("[DONE]")) {
            log.info("OpenAI返回数据结束了");
            session.getBasicRemote().sendText("[DONE]");
            return;
        }
        ObjectMapper objectMapper = new ObjectMapper();
        ChatCompletionResponse response = objectMapper.readValue(data, ChatCompletionResponse.class);
        String reason = response.getChoices().get(0).getFinishReason();
        if (reason.equals("stop") || reason.equals("length")) {
            log.info("OpenAI返回数据结束了");
            session.getBasicRemote().sendText("[DONE]");
            return;
        }
        String content = response.getChoices().get(0).getDelta().getContent();
        log.info("响应内容{}", content);
        session.getBasicRemote().sendText(content);
    }

    @SneakyThrows
    @Override
    public void onFailure(@NotNull EventSource eventSource, @Nullable Throwable t, @Nullable Response response) {
        if (response == null) {
            return;
        }
        ResponseBody body = response.body();
        if (body != null) {
            log.error("OpenAI返回异常数据{}, 异常{}", body.string(), t);
        } else {
            log.error("OpenAI返回异常数据{}, 异常{}", response.toString(), t);
        }
        eventSource.cancel();
    }


    @Override
    public void onClosed(@NotNull EventSource eventSource) {
        log.info("OpenAI关闭连接");
    }

}
